import numpy as np
import pandas as pd
import torch

from config.cfg import Config
from environment.station import FireStation, NewEnergyStation, StorageStation


class Link:
    def __init__(self, internal_exchange, a_out, b_out):
        self.InternalExchangePower = internal_exchange
        self.AOutPower = a_out
        self.BOutPower = b_out


class Area:
    def __init__(self, fire: pd.DataFrame, wind: np.array, wind_capacity, solar: np.array, solar_capacity, load: np.array, storage: pd.DataFrame, water=None):
        self.FireStations = []
        self.WaterStations = []
        self.StorageStations = []
        wind_predict = wind[:, 0].reshape(-1)
        wind_actual = wind[:, 1].reshape(-1)
        solar_predict = solar[:, 0].reshape(-1)
        solar_actual = solar[:, 1].reshape(-1)
        self.Wind = NewEnergyStation(wind_predict, wind_actual, wind_capacity)
        self.Solar = NewEnergyStation(
            solar_predict, solar_actual, solar_capacity)
        self.Load = load
        self.max_load = np.max(load)
        for i in range(len(storage)):
            station = StorageStation(
                node_name=storage['节点名称'].loc[i],
                station_name=storage['厂站名称'].loc[i],
                storage_type=storage['储能类别'].loc[i],
                output=storage['储能容量\n（MW）'].loc[i],
                hours=storage['发电容量小时数\n(小时)'].loc[i],
                hours_in_year=storage['年发电小时设计值\n（小时）'].loc[i],
                rate=storage['充放电效率'].loc[i],
            )
            self.StorageStations.append(station)
        self.StorageStationsSize = len(self.StorageStations)
        self.FireMaxOutput, self.FireMinOutput = 0, 0
        for i in range(len(fire)):
            fireStation = FireStation(
                node_name=fire['节点名称'].loc[i],
                station_name=fire['厂站名称'].loc[i],
                max_output=fire['装机容量（MW）'].loc[i],
                upper_f1=fire['最大出力1(PU.)'].loc[i],
                upper_f2=fire['最大出力2(PU.)'].loc[i],
                low_f1=fire['最小出力1(PU.)'].loc[i],
                low_f2=fire['最小出力2(PU.)'].loc[i],
                restart_cost=fire['启停成本\n（万元）'].loc[i],
                cost_f1=fire['煤耗系数A'].loc[i],
                cost_f2=fire['煤耗系数B'].loc[i],
                cost_f3=fire['煤耗系数C'].loc[i]
            )
            self.FireStations.append(fireStation)
            self.FireMaxOutput += fireStation.OutputRange[3]
            self.FireMinOutput += fireStation.OutputRange[0]
        self.FireStationsSize = len(fire)


class ElectricScheduleEnv:
    def __init__(self, a: Area, b: Area, link: Link, args: Config):
        self.args = args
        self.state_dim = [5, 5]
        self.obs_dim = [5, 5]
        self.max_step = 24*1
        self.A = a
        self.B = b
        self.Areas = [a, b]
        self.Link = link
        self.CurrentTime = 0
        # self.AFireOutputs = [1 for _ in range(a.FireStationsSize)]
        self.AFireOutputs = 1
        # self.BFireOutputs = [1 for _ in range(b.FireStationsSize)]
        self.BFireOutputs = 1
        # total fire + storage
        self.action_dim = [1 + self.A.StorageStationsSize,
                           1 + self.B.StorageStationsSize]
        self.AMaxCapacity = 0
        self.BMaxCapacity = 0
        for s in self.A.StorageStations:
            self.AMaxCapacity += s.MaxCapacity
        for s in self.B.StorageStations:
            self.BMaxCapacity += s.MaxCapacity

    def step(self, actionA, actionB):
        done = False
        aLastFireOutput = self.AFireOutputs
        bLastFireOutput = self.BFireOutputs
        # aWindOutput, bWindOutput = actionA["wind"], actionB['wind']
        # aSolarOutput, bSolarOutput = actionA['solar'], actionB['solar']
        aWindOutput, bWindOutput = self.A.Wind.PredictOutput[
            self.CurrentTime], self.B.Wind.PredictOutput[self.CurrentTime]
        aSolarOutput, bSolarOutput = self.A.Solar.PredictOutput[
            self.CurrentTime], self.B.Solar.PredictOutput[self.CurrentTime]
        self.AFireOutputs, self.BFireOutputs = actionA['fire'], actionB['fire']
        aStorage, bStorage = actionA['storage'], actionB['storage']
        a2b = actionA['exchange']
        buyA = actionA["buy"]
        buyB = actionB["buy"]
        b2a = actionB['exchange']

        rewardA, rewardB = 0, 0
        outputA, outputB = aWindOutput+aSolarOutput, bWindOutput+bSolarOutput
        # 计算风电reward
        # rewardA += self.A.Wind.step(aWindOutput, self.CurrentTime)
        # rewardB += self.B.Wind.step(bWindOutput, self.CurrentTime)
        # 计算太阳能reward
        # rewardA += self.A.Solar.step(aSolarOutput, self.CurrentTime)
        # rewardB += self.B.Solar.step(bSolarOutput, self.CurrentTime)

        # 计算联络线reward
        # to a
        linkCostA, linkCostB = 0, 0
        linkCostA += a2b
        linkCostB += b2a
        # rewardA += linkCostA
        # rewardB += linkCostB
        outputA = outputA - a2b + buyA
        outputB = outputB - b2a + buyB

        # 计算电池reward
        battery_cost = 0
        aCapacity, bCapacity = 0, 0
        for i in range(len(self.A.StorageStations)):
            # 正数放电，负数充电
            outputA += aStorage[i]
            battery_cost += self.A.StorageStations[i].step(aStorage[i])
            aCapacity += self.A.StorageStations[i].Capacity
        rewardA += battery_cost
        battery_cost = 0
        for i in range(len(self.B.StorageStations)):
            outputB += bStorage[i]
            battery_cost += self.B.StorageStations[i].step(bStorage[i])
            bCapacity += self.B.StorageStations[i].Capacity
        rewardB += battery_cost

        # 计算火电reward
        fire_cost = 0
        outputA += self.AFireOutputs
        fire_cost = self.A.FireStations[0].step(
            self.AFireOutputs, aLastFireOutput)
        # for i in range(len(self.A.FireStations)):
        #     outputA += self.AFireOutputs[i]
        #     fire_cost += self.A.FireStations[i].step(
        #         self.AFireOutputs[i],
        #         aLastFireOutput[i],
        #     )
        rewardA += fire_cost

        fire_cost = 0
        outputB += self.BFireOutputs
        fire_cost = self.B.FireStations[0].step(
            self.BFireOutputs, bLastFireOutput)
        # for i in range(len(self.B.FireStations)):
        #     outputB += self.BFireOutputs[i]
        #     fire_cost += self.B.FireStations[i].step(
        #         self.BFireOutputs[i],
        #         bLastFireOutput[i],
        #     )
        rewardB += fire_cost

        # 扣除电能出力和负荷不平衡的成本
        rewardA -= (self.A.Load[self.CurrentTime] - outputA)**2
        # if outputB < self.B.Load[self.CurrentTime]:
        rewardB -= (self.B.Load[self.CurrentTime] - outputB)**2

        self.CurrentTime += 1
        if self.CurrentTime == self.max_step:
            done = True
        return [
            np.array([
                self.CurrentTime,
                self.A.Load[self.CurrentTime] / self.A.max_load,
                self.A.Wind.ActualOutput[self.CurrentTime] /
                self.A.Wind.InstalledCapacity,
                self.A.Solar.ActualOutput[self.CurrentTime] /
                self.A.Solar.InstalledCapacity,
                aCapacity / self.AMaxCapacity
            ], dtype=np.float32),
            np.array([
                self.CurrentTime,
                self.B.Load[self.CurrentTime] / self.B.max_load,
                self.B.Wind.ActualOutput[self.CurrentTime] /
                self.B.Wind.InstalledCapacity,
                self.B.Solar.ActualOutput[self.CurrentTime] /
                self.B.Solar.InstalledCapacity,
                bCapacity / self.BMaxCapacity
            ], dtype=np.float32)
        ], [rewardA*1e-4, rewardB*1e-4], [done, done]

    def reset(self):
        self.CurrentTime = 0
        # reset storage station
        for s in self.A.StorageStations:
            s.Capacity = 0
        for s in self.B.StorageStations:
            s.Capacity = 0
        return [
            np.array([
                0,
                self.A.Load[self.CurrentTime] / self.A.max_load,
                self.A.Wind.PredictOutput[self.CurrentTime] /
                self.A.Wind.InstalledCapacity,
                self.A.Solar.PredictOutput[self.CurrentTime] /
                self.A.Solar.InstalledCapacity,
                0,
            ], dtype=np.float32),
            np.array([
                0,
                self.B.Load[self.CurrentTime] / self.B.max_load,
                self.B.Wind.PredictOutput[self.CurrentTime] /
                self.B.Wind.InstalledCapacity,
                self.B.Solar.PredictOutput[self.CurrentTime] /
                self.B.Solar.InstalledCapacity,
                0,
            ], dtype=np.float32),
        ]


def build_env(root: str, args: Config) -> ElectricScheduleEnv:
    info_fire = pd.read_excel(root+"/data/A/1.样例数据（甲）.xls", sheet_name="热力电厂")
    info_fire = info_fire.groupby('所属区域')
    fireA = info_fire.get_group('A')
    fireA.reset_index(inplace=True, drop=True)
    fireB = info_fire.get_group('B')
    fireB.reset_index(inplace=True, drop=True)

    new_energy_capacity = pd.read_excel(
        root+"/data/A/1.样例数据（甲）.xls", sheet_name="新能源")
    windA_capacity, solarA_capacity = new_energy_capacity[
        '风电装机(MW)'].loc[0], new_energy_capacity['光伏装机(MW)'].loc[0]
    windB_capacity, solarB_capacity = new_energy_capacity[
        '风电装机(MW)'].loc[1], new_energy_capacity['光伏装机(MW)'].loc[1]

    info_new_energy = pd.read_excel(root+"/data/A/2.新能源时序（甲）.xlsx")
    windA = info_new_energy.iloc[2:, 2:4].values
    solarA = info_new_energy.iloc[2:, 4:6].values
    windB = info_new_energy.iloc[2:, 6:8].values
    solarB = info_new_energy.iloc[2:, 8:10].values

    info_load = pd.read_excel(root+"/data/A/1.样例数据（甲）.xls", sheet_name='负荷')
    loadA = info_load['A分区'].loc[1:].values.reshape(-1)
    loadB = info_load['B分区'].loc[1:].values.reshape(-1)

    storage_info = pd.read_excel(root+"/data/A/1.样例数据（甲）.xls", sheet_name='储能')
    storage_info = storage_info.groupby("分区")
    storageA = storage_info.get_group('A')
    storageA.reset_index(inplace=True, drop=True)
    storageB = storage_info.get_group('B')
    storageB.reset_index(inplace=True, drop=True)

    AreaA = Area(
        fire=fireA,
        wind=windA,
        wind_capacity=windA_capacity,
        solar=solarA,
        solar_capacity=solarA_capacity,
        load=loadA,
        storage=storageA
    )
    AreaB = Area(
        fire=fireB,
        wind=windB,
        wind_capacity=windB_capacity,
        solar=solarB,
        solar_capacity=solarB_capacity,
        load=loadB,
        storage=storageB
    )

    info_link = pd.read_excel(root+"/data/A/1.样例数据（甲）.xls", sheet_name="联络线")
    internal = info_link['正向输电能力（MW）'].loc[0]
    a_out = info_link['正向输电能力（MW）'].loc[2]
    b_out = info_link['正向输电能力（MW）'].loc[3]
    link = Link(
        internal_exchange=internal,
        a_out=a_out,
        b_out=b_out
    )
    return ElectricScheduleEnv(AreaA, AreaB, link, args)


env = build_env("..", Config())
